import pickle
import os
import geopandas as gpd
import numpy as np
import multiprocessing as mp
from tqdm import tqdm as tqdm
from common.geometry.geom import Geom
from analytics.crop_recognition.data_utils.location_utils import (
    get_s2_unique_location_id,
)


def build_sample_with_label(src_path, dst_path, label, skip_exist=True):
    if os.path.exists(dst_path):
        if skip_exist:
            return
    with open(src_path, "rb") as f:
        cur_dict = pickle.load(f)
    cur_dict["crop"] = [label for _ in range(len(cur_dict["datetime"]))]

    with open(dst_path, "wb") as f:
        pickle.dump([cur_dict], f)


if __name__ == "__main__":
    gpkg_root = "/NAS6/Members/linchenxi/projects/crop_recognition/dataset/liaoning_area/processed"
    basic_folder = (
        "/NAS6/Members/linchenxi/projects/crop_recognition/new_point_dataset_scl/"
    )
    save_folder = (
        "/NAS6/Members/linchenxi/projects/crop_recognition/points_samples/liaoning"
    )

    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
        os.makedirs(os.path.join(save_folder, "train"))
        os.makedirs(os.path.join(save_folder, "valid"))

    src_gpkg_filelist = []
    for root, _, filenames in os.walk(gpkg_root):
        for filename in filenames:
            if filename.endswith(".gpkg") and "processed" in root:
                # if filename == "Hubei_mamual_rice_2021-2023.gpkg":
                #     continue
                cur_filepath = os.path.join(root, filename)
                src_gpkg_filelist.append(cur_filepath)
    # src_gpkg_filelist = [
    #     # "/NAS6/Members/linchenxi/projects/crop_recognition/dataset/heilongjiang_area/processed/heilongjiang_rice_2023.gpkg", # noqa:E501
    #     # "/NAS6/Members/linchenxi/projects/crop_recognition/dataset/heilongjiang_area/processed/heilongjiang_corn_2023.gpkg", # noqa:E501
    #     "/NAS6/Members/linchenxi/projects/crop_recognition/dataset/heilongjiang_area/processed/henan_trees_all_year.gpkg",  # noqa:E501
    #     "/NAS6/Members/linchenxi/projects/crop_recognition/dataset/heilongjiang_area/processed/henan_urban_all_year.gpkg",  # noqa:E501
    # ]

    for gpkg_path in src_gpkg_filelist:
        print(f"Now processing {gpkg_path}")
        cur_gdf = gpd.read_file(gpkg_path)
        total_len = len(cur_gdf)
        train_flag = np.ones((total_len,), dtype=np.uint8)
        train_flag[: int(total_len * 0)] = 0
        train_flag = np.random.permutation(train_flag)
        pool = mp.Pool(70)
        pbar = tqdm(total=len(cur_gdf))

        def update_pbar(arg=None):
            pbar.update(1)

        for idx, row in cur_gdf.iterrows():
            cur_year = int(row["year"])
            cur_crop = row["crop"]
            cur_year_str = str(cur_year)
            if row.geometry.geom_type == "MultiPoint":
                pbar = tqdm(total=len(row.geometry.geoms))
                for point in row.geometry.geoms:
                    cur_geom = Geom.from_shapely(point)
                    tileid, row, col = get_s2_unique_location_id(cur_geom)
                    cur_filename = f"{tileid}_{row}_{col}_{cur_year_str}.p"
                    dest_filename = f"{tileid}_{row}_{col}_{cur_year_str}_{cur_crop}.p"
                    cur_sample_path = os.path.join(
                        basic_folder, tileid, cur_year_str, cur_filename
                    )
                    if train_flag[idx] == 1:
                        dest_sample_path = os.path.join(
                            save_folder, "train", dest_filename
                        )
                    else:
                        dest_sample_path = os.path.join(
                            save_folder, "valid", dest_filename
                        )
                    # build_sample_with_label(cur_sample_path, dest_sample_path,
                    # cur_crop, True)
                    pool.apply_async(
                        build_sample_with_label,
                        args=[cur_sample_path, dest_sample_path, cur_crop, True],
                        callback=update_pbar,
                        error_callback=print,
                    )
            else:
                cur_geom = Geom.from_shapely(row["geometry"])
                tileid, row, col = get_s2_unique_location_id(cur_geom)
                cur_filename = f"{tileid}_{row}_{col}_{cur_year_str}.p"
                dest_filename = f"{tileid}_{row}_{col}_{cur_year_str}_{cur_crop}.p"
                cur_sample_path = os.path.join(
                    basic_folder, tileid, cur_year_str, cur_filename
                )
                if train_flag[idx] == 1:
                    dest_sample_path = os.path.join(save_folder, "train", dest_filename)
                else:
                    dest_sample_path = os.path.join(save_folder, "valid", dest_filename)
                # build_sample_with_label(cur_sample_path, dest_sample_path,
                # cur_crop, True)
                pool.apply_async(
                    build_sample_with_label,
                    args=[cur_sample_path, dest_sample_path, cur_crop, True],
                    callback=update_pbar,
                    error_callback=print,
                )

        pool.close()
        pool.join()
